# =========================================
#   小易哥的第一个 LLM 微调脚本（讲人话注释版）
# =========================================

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
import torch


# ====== 第1步：加载本地模型路径 ======
# 注意：Windows 路径必须写成 r"xxx" 或者用双斜杠 \\ ，否则会把 \m 当转义字符
model_dir = r"D:\models\qwen\qwen\Qwen3-0___6B"


# ====== 第2步：加载分词器和模型 ======
# tokenizer 负责把文字 -> token（数字），模型训练必须配套使用同款 tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_dir,
    use_fast=False      # 关闭快速分词器，不然有些 Qwen 模型会报错
)

# 加载模型，并让 transformers 自动帮你放到显卡（如果有）
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    device_map="auto",      # 自动选择 GPU / CPU
    dtype=torch.float16     # 使用半精度 FP16，显存直接省一半
)


# ====== 第3步：配置 LoRA 轻量级微调 ======
# LoRA 是一种“只训练很少量参数”的方法，非常适合 1660Super 这种 6G 显卡
# r 越大模型学习能力越强，但显存越高；8 是非常合理的设置
lora_config = LoraConfig(
    r=8,                       # rank，决定可训练矩阵的大小
    lora_alpha=16,             # 缩放系数，通常为2倍rank
    target_modules=["q_proj", "v_proj"],  # 只在注意力层训练，节省显存
    lora_dropout=0.05,         # Dropout 防止过拟合
    bias="none",               # 不训练 bias 项
    task_type="CAUSAL_LM"      # 因果语言模型任务（聊天、文本生成）
)

# 将模型“转成 LoRA 可训练模型”
model = get_peft_model(model, lora_config)


# ====== 第4步：加载 JSONL 数据集 ======
# data.jsonl 的每一行应该是：
# {"instruction": "问题", "output": "回答"}
# 指定文件路径后，自动加载为 HuggingFace Dataset
dataset = load_dataset(
    "json",
    data_files="../data/data_1.jsonl",
    split="train"
)


# ====== 第5步：把 instruction + output 拼成模型训练数据 ======
# 模型训练需要完整的一段“输入 -> 输出”的文本，因此我们把两者合并成 prompt
def tokenize_fn(example):

    # 自己定义的训练格式（随意，只要统一即可）
    # 例如：
    # 用户：你是谁？
    # 小易哥：我是小易哥！
    prompt = f"用户：{example['instruction']}\n小易哥：{example['output']}"

    # 将文本转为 token，设置最大长度防止爆显存
    tokens = tokenizer(
        prompt,
        truncation=True,
        max_length=512
    )

    # labels 就是模型要学习的目标
    # 常规做法是 labels=input_ids 的拷贝
    tokens["labels"] = tokens["input_ids"].copy()

    return tokens

# 将 tokenize_fn 应用到整个数据集
tokenized_dataset = dataset.map(tokenize_fn)


# ====== 第6步：训练参数（适合 1660 Super） ======
training_args = TrainingArguments(
    output_dir="./qwen-finetune-output",  # 输出结果存到这个文件夹
    per_device_train_batch_size=1,        # 每次只喂 1 条，6GB 显存必须这么设
    gradient_accumulation_steps=4,        # 累积 4 次相当于 batch=4（效果更稳定）
    num_train_epochs=3,                   # 训练 3 轮即可，这种小任务 1~3 都可以
    learning_rate=1e-4,                   # 学习率，LoRA 推荐 1e-4 到 2e-4
    fp16=True,                            # 启用半精度训练（更快更省显存）
    logging_steps=1,                      # 每步都打印日志（方便看 loss）
    save_strategy="epoch",                # 每个 epoch 保存一次模型
    save_total_limit=2,                   # 只保留最近 2 个 checkpoint
    report_to="none"                      # 不用 wandb 这些监控平台
)


# ====== 第7步：创建 Trainer（就是帮你跑微调的类） ======
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset
)


# ====== 第8步：开始训练！ ======
trainer.train()
